#include "pch.h"
#include "GRTXScene.h"
#include "GRTXMesh.h"
#include "GRTXBVH.h"
#include "GRTXObject.h"

int MaxPrimitive;
using namespace GMath;
using namespace std;
GRTXScene& GRTXScene::Get()
{
	// TODO: 在此处插入 return 语句
	static GRTXScene* inst = new GRTXScene();
	return *inst;
}
void GRTXScene::InitRTX(std::vector<GUSPtr<GRiObj>>& SceneObjects, const std::unordered_map<std::string, std::unique_ptr<GRiMesh>>& meshes, GUPtr<GRiCamera>& Camera, float width, float height, float PRR,int maxPrimitive)
{
	mMaxPrimitive = maxPrimitive;
	MaxPrimitive = mMaxPrimitive;

	ReBuildMesh(meshes);
	ResetObjs(SceneObjects);

	BuildBVH();
	pCamera = Camera.get();
	Width = width;
	Height = height;
	mPRR = PRR;


}

void GRTXScene::ReBuildMesh(const std::unordered_map<std::string, std::unique_ptr<GRiMesh>>& meshes)
{
	mMeshes.clear();

	for (auto& it : meshes)
	{
		mMeshIdx[it.first] = (int)mMeshes.size();
		mMeshes.emplace_back( new GRTXMesh(it.second->MeshData));
	}
}

void GRTXScene::ResetObjs(std::vector<GUSPtr<GRiObj>>& SceneObjects)
{
	mObjs.clear();
	mObjs.resize(SceneObjects.size());
	GUSPtr<GRTXMesh> mesh;
	GRiMaterial mat;
	GMath::Matrix world;
	GMath::Vector3 emission;
	for (int i=0;i < SceneObjects.size();i ++)
	{
		
		world = SceneObjects[i]->GetComponent<GRiCmpTransform>()->GetWorldMat();

		auto meshId = mMeshIdx[SceneObjects[i]->GetComponent<GRiCmpGeometry>()->GetMesh()->UniqueName];
		mesh = mMeshes[meshId].get();

		mat = SceneObjects[i]->GetComponent<GRiCmpMaterial>()->Mat;

		if (SceneObjects[i]->HasComponent<GRiCmpLight>())
		{
			emission = SceneObjects[i]->GetComponent<GRiCmpLight>()->pLight->Color * SceneObjects[i]->GetComponent<GRiCmpLight>()->pLight->Intensity;
		}
		else
			emission = GMath::Vector3(0.f);

		GSPtr<GRTXObject> temp =make_shared<GRTXObject>(world, mesh, mat, emission);
		mObjs[i] = temp;
	}

	BuildBVH();
}

void GRTXScene::SampleLight(GRiVertex& samplePos, float& pdf) const
{
	// TODO
	float allLightArea = 0;
	for (auto obj:mObjs)
	{
		if (obj->IsLight())
		{
			allLightArea += obj->GetArea();
		}
	}

	allLightArea = allLightArea * GGiEngineUtil::GetRandom();
	for (auto obj : mObjs)
	{
		if (obj->IsLight())
		{
			allLightArea -= obj->GetArea();
			if (allLightArea <= 0)
			{
				obj->SampleLight(samplePos, pdf);
				break;
			}
		}
	}
}

GRTXRay GRTXScene::GenRay(float u, float v)const
{
	GRTXRay ray;
	ray.Orgin = pCamera->GetPosition();
	GMath::Matrix cor = pCamera->GetCornerInWorld(); 
	Vector3 temp0 = cor[0];
	Vector3 temp1 = cor[1];
	Vector3 temp2 = cor[2];
	Vector3 temp3 = cor[3];

	ray.Dir = (temp0 + (temp1 - temp0) * u + (temp2 - temp0) * v).Normalize();

	return ray;
}

GMath::Vector3 GRTXScene::CastRay(const GRTXRay& ray)const
{
	static GMath::Vector3 bg = Vector3(0.0001f);

	auto shaderPoint = pNode->Intersect(ray);


	if (!shaderPoint.isHit)
		return bg;
	//return shaderPoint.p.Normal;
	
	Vector3 L_Dir , L_InDir;

	// 本身是发光项
	if (shaderPoint.isLight)
	{
		L_Dir = shaderPoint.emission;
		return L_Dir;
	}
	
	//////////////////////////////////////////////////////////////////////////
	//// 计算直接光照
	//////////////////////////////////////////////////////////////////////////
	{
		float lightPdf;
		GRiVertex SampleLightPos;

		SampleLight(SampleLightPos,lightPdf);

		// 检测和采样的光源点之中是否存在遮蔽物
		Vector3 o = shaderPoint.p.Position + (shaderPoint.p.Normal) * GGiEngineUtil::EPSILION;
		Vector3 d = (SampleLightPos.Position - o).Normalize();
		GRTXRay toLightRay(o,d);
		auto lightPoint = pNode->Intersect(toLightRay);
		if ((lightPoint.p.Position - SampleLightPos.Position).Length() < GGiEngineUtil::EPSILION)
		{
			L_Dir += lightPoint.emission
				* shaderPoint.mat.BRDF(shaderPoint.p, d, -ray.Dir)
				* shaderPoint.p.Normal.Dot(d)
				* lightPoint.p.Normal.Dot(-d)
				/ ((lightPoint.p.Position - shaderPoint.p.Position).Length() * (lightPoint.p.Position - shaderPoint.p.Position).Length())
				/ lightPdf;
		}
	}
	return L_Dir + L_InDir;
	//////////////////////////////////////////////////////////////////////////
	//// 计算间接光照
	//////////////////////////////////////////////////////////////////////////
	{
		// 俄罗斯轮盘赌成功
		if (GGiEngineUtil::GetRandom() < mPRR)
		{
			GRTXRay randomRay;
			randomRay.Orgin = shaderPoint.p.Position + (shaderPoint.p.Normal) * GGiEngineUtil::EPSILION;	// 避免自相交

			do
			{
				randomRay.Dir = GetRandomRay();
			} while (randomRay.Dir.Dot(shaderPoint.p.Normal) < 0);
		

			auto nexrShaderPoint = pNode->Intersect(randomRay);

			if (nexrShaderPoint.isHit && (!nexrShaderPoint.isLight))
			{
				L_InDir += CastRay(randomRay)
					* shaderPoint.mat.BRDF(shaderPoint.p, randomRay.Dir, -ray.Dir)
					* shaderPoint.p.Normal.Dot(randomRay.Dir)
					/ shaderPoint.mat.PDF(shaderPoint.p, randomRay.Dir, -ray.Dir)
					/ mPRR;
			}
		}
	}
	

	return L_Dir + L_InDir;

}

void GRTXScene::BuildBVH()
{
	std::vector<GSPtr<GRTXBVH>> bvhNodes;
	for (auto& obj : mObjs)
	{
		GSPtr<GRTXBVH> node = make_shared<GRTXBVH>();

		node->Box = obj->GetBox();
		node->Left = obj->pMesh->pBVHNode;
		node->Right = nullptr;
		node->Type = GRTXBVH::NodeType::NTInstance;
		node->pObj = obj;

		bvhNodes.push_back(node);
	}
	
	pNode = move(RecursiveBuild(bvhNodes));
}

GSPtr<GRTXBVH> GRTXScene::RecursiveBuild(std::vector<GSPtr<GRTXBVH>> objs)
{
	GSPtr<GRTXBVH> node(new GRTXBVH());
 	for (auto& it : objs)
 	{
 		node->Box.Union(it->Box);
 	}
 
 	if (objs.size() == 1)
 	{
 		node = objs[0];
 	}
 	else if (objs.size() == 2)
 	{
		node->Left = RecursiveBuild({ objs[0]});
		node->Right = RecursiveBuild({objs[1]});
 	}
 	else
 	{
 		int dim = node->Box.MaxExtent();
 		sort(objs.begin(), objs.end(),
 			[dim](GSPtr<GRTXBVH>& n1, auto& n2) {
 				return n1->Box.Center()[dim] < n2->Box.Center()[dim];
 			});
	
		auto mid = objs.begin() + objs.size() / 2;
		node->Left = move(RecursiveBuild({ objs.begin(),mid }));
		node->Right = move(RecursiveBuild({ mid,objs.end() }));
 	}

	return move(node);
}
